#pragma once

#include <memory>
#include <vector>

#include "Error.h"
#include "TypeDef.h"
#include "VisitorInterface.h"

namespace FlooP
{
	// Note all AST nodes are immutable and the constructors have move semantics, so be careful!
	class AST;
	class ASTProcedure;
	class ASTStatement;
	class ASTExpression;
	class ASTProcedureCall;

	class AST
	{
		public:

		AST (std::vector<std::shared_ptr<ASTProcedure>>& procedure, std::vector<std::shared_ptr<ASTProcedureCall>>& calls)
		{
			m_procedure.swap (procedure);
		}

		void accept (VisitorInterface& visitor) const
		{
			visitor.visit (*this);
		}

		const std::vector<std::shared_ptr<ASTProcedure>>& procedure () const
		{
			return m_procedure;
		}

		const std::vector<std::shared_ptr<ASTProcedureCall>>& calls () const
		{
			return m_calls;
		}

		private:

		std::vector<std::shared_ptr<ASTProcedure>>	   m_procedure;
		std::vector<std::shared_ptr<ASTProcedureCall>> m_calls;
	};

	class ASTProcedure
	{
		public:

		ASTProcedure (std::string& name, std::vector<std::string>& args, std::shared_ptr<ASTStatement> body)
		{
			m_name.swap (name);
			m_args.swap (args);
			m_body = body;
		}

		const std::string& name () const
		{
			return m_name;
		}

		const std::vector<std::string>& args () const
		{
			return m_args;
		}

		const std::shared_ptr<ASTStatement>& body () const
		{
			return m_body;
		}

		void accept (VisitorInterface& visitor) const
		{
			visitor.visit (*this);
		}

		private:

		std::string m_name;

		std::vector<std::string> m_args;
		std::shared_ptr<ASTStatement> m_body;
	};

	class ASTStatement
	{
		public:

		virtual ~ASTStatement () {}

		virtual void accept (VisitorInterface& visitor, std::shared_ptr<Procedure> func) const = 0;
	};

	class ASTBlock : public ASTStatement
	{
		public:

		ASTBlock (std::vector<std::shared_ptr<ASTStatement>>& block, uint32_t block_num)
		:
			m_block_num (block_num)
		{
			m_block.swap (block);
		}

		uint32_t block_num () const
		{
			return m_block_num;
		}

		const std::vector<std::shared_ptr<ASTStatement>>& stmts () const
		{
			return m_block;
		}

		void accept (VisitorInterface& visitor, std::shared_ptr<Procedure> func) const
		{
			visitor.visit (*this, func);
		}

		private:

		uint32_t m_block_num;
		std::vector<std::shared_ptr<ASTStatement>> m_block;
	};

	class ASTIf : public ASTStatement
	{
		public:

		ASTIf (std::shared_ptr<ASTExpression> cond, std::shared_ptr<ASTStatement> if_part)
		:
			m_cond	  (cond),
			m_if_part (if_part)
		{}

		std::shared_ptr<ASTExpression> cond () const
		{
			return m_cond;
		}

		std::shared_ptr<ASTStatement> if_part () const
		{
			return m_if_part;
		}

		void accept (VisitorInterface& visitor, std::shared_ptr<Procedure> func) const
		{
			visitor.visit (*this, func);
		}
		private:

		std::shared_ptr<ASTExpression> m_cond;
		std::shared_ptr<ASTStatement>  m_if_part;
	};

	class ASTLoop : public ASTStatement
	{
		public:

		ASTLoop (std::shared_ptr<ASTExpression> times, std::shared_ptr<ASTStatement> body)
		:
			m_times (times),
			m_body  (body)
		{}

		std::shared_ptr<ASTExpression> times () const
		{
			return m_times;
		}

		std::shared_ptr<ASTStatement> body () const
		{
			return m_body;
		}

		void accept (VisitorInterface& visitor, std::shared_ptr<Procedure> func) const
		{
			visitor.visit (*this, func);
		}
		private:

		std::shared_ptr<ASTExpression> m_times;
		std::shared_ptr<ASTStatement>  m_body;
	};

	class ASTMuLoop : public ASTStatement
	{
		public:

		ASTMuLoop (std::shared_ptr<ASTStatement> body)
		:
			m_body  (body)
		{}


		std::shared_ptr<ASTStatement> body () const
		{
			return m_body;
		}

		void accept (VisitorInterface& visitor, std::shared_ptr<Procedure> func) const
		{
			visitor.visit (*this, func);
		}
		private:

		std::shared_ptr<ASTStatement>  m_body;
	};

	class ASTAbort : public ASTStatement
	{
		public:

		ASTAbort (uint32_t block)
		:
			m_block (block)
		{}

		uint32_t block () const
		{
			return m_block;
		}

		void accept (VisitorInterface& visitor, std::shared_ptr<Procedure> func) const
		{
			visitor.visit (*this, func);
		}

		private:

		uint32_t m_block;
	};

	class ASTQuit : public ASTStatement
	{
		public:

		ASTQuit (uint32_t block)
		:
			m_block (block)
		{}

		uint32_t block () const
		{
			return m_block;
		}

		void accept (VisitorInterface& visitor, std::shared_ptr<Procedure> func) const
		{
			visitor.visit (*this, func);
		}

		private:

		uint32_t m_block;
	};

	class ASTAssignment : public ASTStatement
	{
		public:

		ASTAssignment (uint32_t cell, std::shared_ptr<ASTExpression> expr)
		:
			m_cell (cell),
			m_expr (expr)
		{}

		uint32_t cell () const
		{
			return m_cell;
		}

		std::shared_ptr<ASTExpression> expr () const
		{
			return m_expr;
		}

		void accept (VisitorInterface& visitor, std::shared_ptr<Procedure> func) const
		{
			visitor.visit (*this, func);
		}

		private:

		uint32_t m_cell;
		std::shared_ptr<ASTExpression> m_expr;
	};

	class ASTExpression
	{
		public:

		enum Type
		{
			COMPARE,
			BINARY_OP,
			CELL,
			OUTPUT,
			CALL,
			IDENTIFIER,
			NUMBER,
			STRING
		};

		virtual ~ASTExpression () {}

		virtual Type type () const = NULL;

		virtual void accept (VisitorInterface& visitor, std::shared_ptr<Procedure> func) const = 0;
	};

	class ASTCompare : public ASTExpression
	{
		public:

		enum OpType
		{
			OR, 
			AND, 
			EQUALS, 
			LESS_THAN,  
			GREATER_THAN,
		};

		ASTCompare (OpType type, std::shared_ptr<ASTExpression> left, std::shared_ptr<ASTExpression> right)
		:
			m_left  (left),
			m_right (right),
			m_type  (type)
		{}

		std::shared_ptr<ASTExpression> left () const
		{
			return m_left;
		}

		std::shared_ptr<ASTExpression> right () const
		{
			return m_right;
		}

		ASTExpression::Type type () const
		{
			return ASTExpression::BINARY_OP;
		}

		ASTCompare::OpType op_type () const
		{
			return m_type;
		}

		void accept (VisitorInterface& visitor, std::shared_ptr<Procedure> func) const
		{
			visitor.visit (*this, func);
		}

		private:

		OpType m_type;

		std::shared_ptr<ASTExpression> m_left;
		std::shared_ptr<ASTExpression> m_right;
	};

	class ASTBinaryMathOp : public ASTExpression
	{
		public:

		enum OpType
		{
			PLUS,
			STAR
		};

		class ASTBinaryMathOp (OpType type, std::shared_ptr<ASTExpression> left, std::shared_ptr<ASTExpression> right)
		:
			m_left  (left),
			m_right (right),
			m_type  (type)
		{}

		std::shared_ptr<ASTExpression> left () const
		{
			return m_left;
		}

		std::shared_ptr<ASTExpression> right () const
		{
			return m_right;
		}

		ASTExpression::Type type () const
		{
			return ASTExpression::BINARY_OP;
		}

		ASTBinaryMathOp::OpType op_type () const
		{
			return m_type;
		}

		void accept (VisitorInterface& visitor, std::shared_ptr<Procedure> func) const
		{
			visitor.visit (*this, func);
		}

		private:

		OpType m_type;

		std::shared_ptr<ASTExpression> m_left;
		std::shared_ptr<ASTExpression> m_right;
	};

	class ASTCell : public ASTExpression
	{
		public:

		ASTCell (uint32_t cell)
		:
			m_cell (cell)
		{}

		uint32_t value () const
		{
			return m_cell;
		}

		Type type () const
		{
			return ASTExpression::CELL;
		}

		void accept (VisitorInterface& visitor, std::shared_ptr<Procedure> func) const
		{
			visitor.visit (*this, func);
		}

		private:

		uint32_t m_cell;
	};

	class ASTProcedureCall : public ASTExpression
	{
		public:

		ASTProcedureCall (std::string& name, std::vector<std::shared_ptr<ASTExpression>>& args)
		{
			m_name.swap (name);
			m_args.swap (args);
		}

		const std::string& name () const
		{
			return m_name;
		}

		const std::vector<std::shared_ptr<ASTExpression>>& args () const
		{
			return m_args;
		}

		Type type () const
		{
			return ASTExpression::CALL;
		}

		void accept (VisitorInterface& visitor, std::shared_ptr<Procedure> func) const
		{
			visitor.visit (*this, func);
		}

		private:

		Type m_type;

		std::string m_name;
		std::vector<std::shared_ptr<ASTExpression>> m_args;
	};

	class ASTIdentifier : public ASTExpression
	{
		public:

		ASTIdentifier (std::string& name)
		{
			m_name.swap (name);
		}

		ASTIdentifier (std::string&& name)
		{
			m_name.swap (name);
		}

		Type type () const
		{
			return ASTExpression::IDENTIFIER;
		}

		std::string name () const
		{
			return m_name;
		}

		void accept (VisitorInterface& visitor, std::shared_ptr<Procedure> func) const
		{
			visitor.visit (*this, func);
		}

		private:

		std::string m_name;
	};

	class ASTNumber : public ASTExpression
	{
		public:

		ASTNumber (uint32_t val)
		:
			m_val (val)
		{}

		uint32_t value () const
		{
			return m_val;
		}

		Type type () const
		{
			return ASTExpression::NUMBER;
		}

		void accept (VisitorInterface& visitor, std::shared_ptr<Procedure> func) const
		{
			visitor.visit (*this, func);
		}

		private:

		uint32_t m_val;
	};

	class ASTString : public ASTExpression
	{
		public:

		ASTString (std::string& text)
		{
			m_text.swap (text);
		}

		ASTString (std::string&& text)
		{
			m_text.swap (text);
		}

		Type type () const
		{
			return ASTExpression::STRING;
		}

		const std::string& text () const
		{
			return m_text;
		}

		void accept (VisitorInterface& visitor, std::shared_ptr<Procedure> func) const
		{
			visitor.visit (*this, func);
		}
		private:

		std::string m_text;
	};
}